package yang.yu.sunny.infra;

import org.dbunit.DatabaseUnitException;
import org.dbunit.database.DatabaseConnection;
import org.dbunit.database.DatabaseSequenceFilter;
import org.dbunit.database.IDatabaseConnection;
import org.dbunit.dataset.DataSetException;
import org.dbunit.dataset.FilteredDataSet;
import org.dbunit.dataset.IDataSet;
import org.dbunit.dataset.xml.FlatDtdDataSet;
import org.dbunit.dataset.xml.FlatXmlDataSet;
import org.dbunit.dataset.xml.FlatXmlDataSetBuilder;
import org.dbunit.operation.DatabaseOperation;
import org.xml.sax.InputSource;

import javax.sql.DataSource;
import java.io.*;
import java.sql.SQLException;

/**
 * Created by yyang on 15/10/16.
 */
public class DbUnit {

    private DataSource dataSource;

    private String dtdFile = "/dataset.dtd";

    private String flatXmlDataFile = "/default-data.xml";



    public DbUnit(DataSource dataSource) {
        this.dataSource = dataSource;
    }

    /**
     * 将数据库中的数据导出到XML文件中。
     */
    public void exportData() throws IOException {
        IDatabaseConnection connection = null;
        try {
            connection = new DatabaseConnection(dataSource.getConnection());
            IDataSet dataSet = new FilteredDataSet(new DatabaseSequenceFilter(connection),
                    connection.createDataSet());
            OutputStream out = new FileOutputStream(getClass().getResource(flatXmlDataFile).getFile());
            FlatXmlDataSet.write(dataSet, out);
        } catch (DatabaseUnitException e) {
            e.printStackTrace();
        } catch (SQLException e) {
            e.printStackTrace();
        } finally {
            try {
                connection.close();
            } catch (SQLException e) {
                e.printStackTrace();
            }
        }
    }

    /**
     * 读取数据库表结构，生成XML数据文件的DTD文件
     */
    public void exportDtd() throws IOException {
        IDatabaseConnection connection = null;
        try {
            connection = new DatabaseConnection(dataSource.getConnection());
            IDataSet dataSet = new FilteredDataSet(new DatabaseSequenceFilter(connection),
                    connection.createDataSet());
            Writer out = new OutputStreamWriter(new FileOutputStream(getClass().getResource(dtdFile).getFile()));
            FlatDtdDataSet.write(dataSet, out);
        } catch (DatabaseUnitException e) {
            e.printStackTrace();
        } catch (SQLException e) {
            e.printStackTrace();
        } finally {
            try {
                connection.close();
            } catch (SQLException e) {
                e.printStackTrace();
            }
        }
    }

    public void importData() {
        IDatabaseConnection connection = null;
        try {
            connection = new DatabaseConnection(dataSource.getConnection());
            DatabaseOperation.CLEAN_INSERT.execute(connection, getDatasetFromFile(flatXmlDataFile));
        } catch (DatabaseUnitException e) {
            e.printStackTrace();
        } catch (SQLException e) {
            e.printStackTrace();
        } finally {
            try {
                connection.close();
            } catch (SQLException e) {
                e.printStackTrace();
            }
        }
    }

    private IDataSet getDatasetFromFile(String flatXmlDataFile) {
        return getDatasetFromInputStream(getClass().getResourceAsStream(flatXmlDataFile));
    }

    private IDataSet getDatasetFromInputStream(InputStream in) {
        FlatXmlDataSetBuilder builder = new FlatXmlDataSetBuilder();
        builder.setColumnSensing(true);
        try {
            return builder.build(new InputSource(in));
        } catch (DataSetException e) {
            e.printStackTrace();
            throw new RuntimeException("Cannot get dataset.", e);
        }
        /*IDataSetProducer producer = new FlatXmlProducer(new InputSource(in), false);
        try {
            return new CachedDataSet(producer);
        } catch (DataSetException e) {
            throw new RuntimeException("Cannot get dataset.", e);
        }*/
    }
}
